{
  "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Acrobot\n",
    "\n",
    "*We recommend you look at the [Introduction to Crocoddyl](introduction_to_crocoddyl.ipynb) example before this one.*\n",
    "\n",
    "In the example, we model the acrobot control problem using Crocoddyl. An acrobot is a two joint planar robot with only one actuator. It is a canonnical example of an underactuated system and so presents an interesting control problem.\n",
    "\n",
    "We demonstrate how to:\n",
    "\n",
    "1. Load a model from an urdf.\n",
    "1. Define an actuation mapping for the system.\n",
    "1. Construct and solve the control problem.\n",
    "\n",
    "## Loading the model\n",
    "A standalone double pendulum robot urdf is provided in the [example-robot-data](https://github.com/Gepetto/example-robot-data) repository, this comes bundled with Crocoddyl. Let's load the model and inspect its properties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import pathlib\n",
    "import crocoddyl\n",
    "import pinocchio\n",
    "\n",
    "# Get the path to the urdf\n",
    "from example_robot_data.path import EXAMPLE_ROBOT_DATA_MODEL_DIR\n",
    "\n",
    "urdf_model_path = pathlib.Path(\n",
    "    \"double_pendulum_description\", \"urdf\", \"double_pendulum_simple.urdf\"\n",
    ")\n",
    "urdf_model_path = os.path.join(EXAMPLE_ROBOT_DATA_MODEL_DIR, urdf_model_path)\n",
    "\n",
    "# Now load the model (using pinocchio)\n",
    "robot = pinocchio.robot_wrapper.RobotWrapper.BuildFromURDF(str(urdf_model_path))\n",
    "\n",
    "# The model loaded from urdf (via pinicchio)\n",
    "print(robot.model)\n",
    "\n",
    "# Create a multibody state from the pinocchio model.\n",
    "state = crocoddyl.StateMultibody(robot.model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`pinocchio` comes with some handy wrappers that load a robot along with vizual and collision models. These are all defined via that urdf. `robot.model` is the model is a model of the DAEs (Differential Algebraic Equations).\n",
    "\n",
    "You will notice that the there are two joint configurations `nq` and velocities `nv`.\n",
    "\n",
    "## Actuation Mapping\n",
    "In order to create an underactuated double pendulum, the acrobot, we will create mapping between control inputs and joint torques. This is done by inheriting from `ActuationModelAbstract`. See also `ActuationModelFloatingBase` and `ActuationModelFull` for other options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the control signal to actuated joint mapping\n",
    "class AcrobotActuationModel(crocoddyl.ActuationModelAbstract):\n",
    "    def __init__(self, state):\n",
    "        nu = 1  # Control dimension\n",
    "        crocoddyl.ActuationModelAbstract.__init__(self, state, nu=nu)\n",
    "\n",
    "    def calc(self, data, x, u):\n",
    "        assert len(data.tau) == 2\n",
    "        # Map the control dimensions to the joint torque\n",
    "        data.tau[0] = 0\n",
    "        data.tau[1] = u\n",
    "\n",
    "    def calcDiff(self, data, x, u):\n",
    "        # Specify the actuation jacobian\n",
    "        data.dtau_du[0] = 0\n",
    "        data.dtau_du[1] = 1\n",
    "\n",
    "\n",
    "# Also see ActuationModelFloatingBase and ActuationModelFull\n",
    "actuationModel = AcrobotActuationModel(state)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constructing the Problem\n",
    "\n",
    "Before we solve the control problem, we need to construct the cost models and action models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dt = 1e-3  # Time step\n",
    "T = 1000  # Number of knots\n",
    "\n",
    "# Cost models\n",
    "runningCostModel = crocoddyl.CostModelSum(state, nu=actuationModel.nu)\n",
    "terminalCostModel = crocoddyl.CostModelSum(state, nu=actuationModel.nu)\n",
    "\n",
    "# Add a cost for the configuration positions and velocities\n",
    "xref = np.array([0, 0, 0, 0])  # Desired state\n",
    "stateResidual = crocoddyl.ResidualModelState(state, xref=xref, nu=actuationModel.nu)\n",
    "stateCostModel = crocoddyl.CostModelResidual(state, stateResidual)\n",
    "runningCostModel.addCost(\"state_cost\", cost=stateCostModel, weight=1e-5 / dt)\n",
    "terminalCostModel.addCost(\"state_cost\", cost=stateCostModel, weight=1000)\n",
    "\n",
    "# Add a cost on control\n",
    "controlResidual = crocoddyl.ResidualModelControl(state, nu=actuationModel.nu)\n",
    "bounds = crocoddyl.ActivationBounds(np.array([-1.0]), np.array([1.0]))\n",
    "activation = crocoddyl.ActivationModelQuadraticBarrier(bounds)\n",
    "controlCost = crocoddyl.CostModelResidual(\n",
    "    state, activation=activation, residual=controlResidual\n",
    ")\n",
    "runningCostModel.addCost(\"control_cost\", cost=controlCost, weight=1e-1 / dt)\n",
    "\n",
    "# Create the action models for the state\n",
    "runningModel = crocoddyl.IntegratedActionModelEuler(\n",
    "    crocoddyl.DifferentialActionModelFreeFwdDynamics(\n",
    "        state, actuationModel, runningCostModel\n",
    "    ),\n",
    "    dt,\n",
    ")\n",
    "terminalModel = crocoddyl.IntegratedActionModelEuler(\n",
    "    crocoddyl.DifferentialActionModelFreeFwdDynamics(\n",
    "        state, actuationModel, terminalCostModel\n",
    "    ),\n",
    "    0.0,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we define the control problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a shooting problem\n",
    "q0 = np.zeros((state.nq,))  # Inital joint configurations\n",
    "q0[0] = np.pi / 2  # Down\n",
    "v0 = np.zeros((state.nv,))  # Initial joint velocities\n",
    "x0 = np.concatenate((q0, v0))  # Inital robot state\n",
    "problem = crocoddyl.ShootingProblem(x0, [runningModel] * T, terminalModel)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test the system with a rollout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the problem with a rollout\n",
    "us = [0.01 * np.ones((1,))] * T\n",
    "xs = problem.rollout(us)\n",
    "\n",
    "# Handy to blat up the state and control trajectories\n",
    "crocoddyl.plotOCSolution(xs, us, show=False, figIndex=99, figTitle=\"Test rollout\")\n",
    "\n",
    "# Put a grid on the plots\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig = plt.gcf()\n",
    "axs = fig.axes\n",
    "for ax in axs:\n",
    "    ax.grid()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can solve the optimal control problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now stabilize the acrobot using FDDP\n",
    "solver = crocoddyl.SolverFDDP(problem)\n",
    "\n",
    "# Solve\n",
    "callbacks = []\n",
    "callbacks.append(crocoddyl.CallbackLogger())\n",
    "callbacks.append(crocoddyl.CallbackVerbose())\n",
    "solver.setCallbacks(callbacks)\n",
    "solver.solve([], [], 300, False, 1e-5)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the trajectory with `meshcat` or using `gepetto-gui` (you will need to install [gepetto-viewer]() and [gepetto-viewer-corba]() and start the process in a separate terminal.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display using meshcat\n",
    "robot_display = crocoddyl.MeshcatDisplay(robot, -1, 1, False)\n",
    "display(robot_display.robot.viewer.jupyter_cell())\n",
    "robot_display.displayFromSolver(solver)\n",
    "\n",
    "# Display using gepetto-gui\n",
    "if False:\n",
    "    robot_display = crocoddyl.GepettoDisplay(robot, floor=False)\n",
    "    robot_display.displayFromSolver(solver)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot the trajectory and the solver's convergence properties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotting the solution and the DDP convergence\n",
    "log = solver.getCallbacks()[0]\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "crocoddyl.plotOCSolution(\n",
    "    xs=log.xs, us=log.us, show=False, figIndex=1, figTitle=\"Solution\"\n",
    ")\n",
    "fig = plt.gcf()\n",
    "axs = fig.axes\n",
    "for ax in axs:\n",
    "    ax.grid(True)\n",
    "\n",
    "crocoddyl.plotConvergence(\n",
    "    log.costs,\n",
    "    log.pregs,\n",
    "    log.dregs,\n",
    "    log.grads,\n",
    "    log.stops,\n",
    "    log.steps,\n",
    "    show=False,\n",
    "    figIndex=2,\n",
    ")\n",
    "fig = plt.gcf()\n",
    "axs = fig.axes\n",
    "for ax in axs:\n",
    "    ax.grid(True)\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
